from unittest import TestCase

from comp_arch.loops import loop_case_1, loop_case_1a, loop_case_2, loop_case_2a

from .data import LOOP_CASE_1_A, LOOP_CASE_1_B


class TestLoops(TestCase):
    def test_loop_case_1(self):
        A = [i for i in range(100)]
        B = [i for i in range(101)]
        C = [i for i in range(100)]
        D = [i for i in range(100)]
        expected_a = LOOP_CASE_1_A
        expected_b = LOOP_CASE_1_B
        result_a, result_b = loop_case_1(A, B, C, D)
        self.assertEqual(result_a, expected_a)
        self.assertEqual(result_b, expected_b)

    def test_loop_case_1a(self):
        A = [i for i in range(100)]
        B = [i for i in range(101)]
        C = [i for i in range(100)]
        D = [i for i in range(100)]
        expected_a = LOOP_CASE_1_A
        expected_b = LOOP_CASE_1_B
        result_a, result_b = loop_case_1a(A, B, C, D)
        self.assertEqual(result_a, expected_a)
        self.assertEqual(result_b, expected_b)

    def test_loop_case_2(self):
        A = [[i for i in range(100)] for j in range(100)]
        B = [i for i in range(100)]
        C = [i for i in range(100)]
        result = loop_case_2(A, B, C)
        result_a = loop_case_2a(A, B, C)
        self.assertEqual(result, result_a)
